package inference;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.SortedSet;

import utils.ArrayUtils;
import utils.BackboneCreator;
import utils.IntPair;
import utils.Pair;
import utils.PrettyPrinter;
import model.ModelParameters;
import model.TableInformation;
import model.TableInformationReader;

import com.analog.lyric.dimple.model.Discrete;
import com.analog.lyric.dimple.model.Factor;
import com.analog.lyric.dimple.factorfunctions.core.FactorTable;
import com.analog.lyric.dimple.factorfunctions.core.IFactorTable;
import com.analog.lyric.dimple.model.FactorGraph;
import com.analog.lyric.dimple.schedulers.RandomWithReplacementScheduler;
import com.analog.lyric.dimple.schedulers.RandomWithoutReplacementScheduler;
import com.analog.lyric.dimple.schedulers.TreeOrSequentialScheduler;
import com.analog.lyric.dimple.schedulers.schedule.ISchedule;
import com.analog.lyric.dimple.schedulers.scheduleEntry.IScheduleEntry;
import com.analog.lyric.dimple.schedulers.scheduleEntry.NodeScheduleEntry;
import com.analog.lyric.dimple.solvers.gibbs.Solver;
import com.analog.lyric.dimple.solvers.interfaces.ISolverFactorGraph;
import com.analog.lyric.dimple.solvers.sumproduct.SFactorGraph;
import com.analog.lyric.dimple.solvers.sumproduct.STableFactor;

public class TableInferer {
	private static boolean VERBOSE = false;
	
	private FactorGraph factor_graph;
	private ModelParameters model_params;
	private int[] set_backbone = null;
	
	private int[] best_backbone;
	private int best_backbone_length;
	private double best_bethe_free_energy;
	private double[][][] best_entity_vars_beliefs;
	private double[][] best_predicates_vars_beliefs;
	private double[][] best_col_type_vars_beliefs;
	private double[] best_topic_var_beliefs;
	private double[] best_topic_type_var_beliefs;
	private int best_topic_var_idx = -1;
	private int[] prev_backbone;
	private ArrayList<Pair<int[], Double> > backbones_checked;
	
	private Factor[][] current_entity_factors;
	private Factor[] current_predicate_factors;
	private Factor current_topic_factor;
	private HashSet<Factor> current_col_type_factors;
	
	private boolean use_best_backbone = true;
	
	public ModelParameters getModelParameters() {
		return model_params;
	}
	
	public int[] getBestBackbone() {
		return best_backbone;
	}
	
	public double getBestBetheFreeEnergy() {
		return best_bethe_free_energy;
	}
	
	public ArrayList<Pair<int[], Double> > getBackbonesChecked() {
		return backbones_checked;
	}
	
	public Pair<String[], double[]> getSortedTopics() {
		String[] sorted_output_topic_domain = ArrayUtils.translateIdArrayToStringArray(
				model_params.getTopicVariable(best_topic_var_idx).getDomain().getElements(),
				model_params.getTableInformation());
		Pair<String[], double[]> sorted_pair = ArrayUtils.sortPairedArrays(sorted_output_topic_domain, best_topic_var_beliefs);
		return sorted_pair;
	}
	
	public Pair<String[], double[]> getSortedPredicatesForColumn(int col) {
		String[] sorted_output_pred_domain = ArrayUtils.translateIdArrayToStringArray(
				model_params.getPredicateVariable(col).getDomain().getElements(),
				model_params.getTableInformation());
		Pair<String[], double[]> sorted_pair = ArrayUtils.sortPairedArrays(sorted_output_pred_domain,
				best_predicates_vars_beliefs[col]);
		return sorted_pair;
	}
	
	public Pair<String[], double[]> getSortedEntitiesForCell(int row, int col) {
		String[] sorted_output_entity_domain = ArrayUtils.translateIdArrayToStringArray(
				model_params.getEntityVariable(row, col).getDomain().getElements(),
				model_params.getTableInformation());
		Pair<String[], double[]> sorted_pair = ArrayUtils.sortPairedArrays(sorted_output_entity_domain,
				best_entity_vars_beliefs[row][col]);
		return sorted_pair;
	}
	
	public int getBestBackboneLength() {
		return best_backbone_length;
	}
	
	public void setup(String entity_filename, String type_filename, String fact_filename,
			String predicate_info_filename, String predicate_to_mediator_filename, String fixed_backbone) throws IOException {
		TableInformation table_info = TableInformationReader.readTableInformation(entity_filename, type_filename,
				fact_filename, predicate_info_filename, predicate_to_mediator_filename);

		model_params = new ModelParameters(table_info);
		model_params.createVariablesForModel();
		model_params.recreateHiddenBackboneAdjacencyMatrix();
		model_params.createPredicateFactorCache();
		model_params.createCellFactorCache();
		model_params.createTopicFactorCache();
		
		if (!fixed_backbone.equals("")) {
			String[] backbone_str = fixed_backbone.split(",");
			set_backbone = new int[backbone_str.length];
			for (int c = 0; c < backbone_str.length; c++) {
				set_backbone[c] = Integer.parseInt(backbone_str[c]);
			}
		}
		
		factor_graph = new FactorGraph();
//		factor_graph.setSolverFactory(new com.analog.lyric.dimple.solvers.gaussian.Solver());
//		factor_graph.setSolverFactory(new com.analog.lyric.dimple.solvers.gibbs.Solver());
//		factor_graph.setScheduler(new com.analog.lyric.dimple.schedulers.RandomWithReplacementScheduler());

		current_col_type_factors = new HashSet<Factor>();
		
		this.best_backbone = null;
		this.best_entity_vars_beliefs = new double[table_info.numRows()][model_params.getMaxColumns()][];
		this.best_predicates_vars_beliefs = new double[model_params.getMaxColumns()][];
		this.best_col_type_vars_beliefs = new double[model_params.getMaxColumns()][];
		this.best_topic_var_beliefs = null;
		this.best_topic_type_var_beliefs = null;
		this.prev_backbone = null;
		
		this.current_entity_factors = new Factor[table_info.numRows()][model_params.getMaxColumns()];
		this.current_predicate_factors = new Factor[model_params.getMaxColumns()];
		this.current_topic_factor = null;
	}
	
	private double runInferenceForBackbone(int[] backbone) {
		TableInformation table_info = model_params.getTableInformation();
		int max_num_cols = model_params.getMaxColumns();
		int root_col = ArrayUtils.indexOf(backbone, -1);
		for (Factor col_type_factor : current_col_type_factors) {
			factor_graph.remove(col_type_factor);
		}
		current_col_type_factors.clear();
		for (int col = 0; col < table_info.numVisibleCols(); col++) {
			Discrete pred_var = model_params.getPredicateVariable(col);
			Discrete col_type_var = model_params.getColumnTypeVariable(col);

			Pair<int[][], double[]> sparse_info = ColumnTypeFactor.createFactor(pred_var, col_type_var, model_params);
			IFactorTable column_type_table = FactorTable.create(sparse_info.getFirst(), sparse_info.getSecond(),
					pred_var.getDomain(), col_type_var.getDomain());
//			PrettyPrinter.prettyPrintPredicateFactor(column_type_table, pred_var, 
//					col_type_var, table_info);
			Factor column_type_factor = factor_graph.addFactor(column_type_table, pred_var, col_type_var);
			column_type_factor.setDirectedTo(col_type_var);
			column_type_factor.setLabel("ColType" + col);
//			((STableFactor)column_type_factor.getSolver()).setK(5);
			current_col_type_factors.add(column_type_factor);
		}
		for (int i = table_info.numVisibleCols(); i < backbone.length; i++) {
			if (backbone[i] == -2) {
				break;
			}
			Factor column_type_factor = setUpColumnTypeFactor(i);
			current_col_type_factors.add(column_type_factor);
		}
		if (model_params.ignoreTableTopic()) {
			
		} else {
			int col = root_col;
			Factor prev_factor = current_topic_factor;
			if (prev_factor != null) {
				factor_graph.remove(prev_factor);
			}
			IFactorTable topic_table = model_params.getTopicFactor(col);
			Factor topic_factor = factor_graph.addFactor(topic_table, model_params.getTopicTypeVariable(col),
					model_params.getTopicVariable(col));
			topic_factor.setDirectedTo(model_params.getTopicVariable(col));
			topic_factor.setLabel("Topic" + col);
//			STableFactor sfactor = (STableFactor)(topic_factor.getSolver());
//			sfactor.setK(5);
//			Factor topic_factor = factor_graph.addFactor(topic_table, model_params.getTopicVariable(col),
//					model_params.getTopicTypeVariable(col));
//			topic_factor.setDirectedTo(model_params.getTopicTypeVariable(col));
			current_topic_factor = topic_factor;
			
			prev_factor = current_predicate_factors[col];
			if (prev_factor != null) {
				factor_graph.remove(prev_factor);
			}
			IFactorTable pred_table = model_params.getPredicateFactor(max_num_cols, col);
//			PrettyPrinter.prettyPrintPredicateFactor(pred_table, model_params.getTopicVariable(col), 
//					model_params.getPredicateVariable(col), table_info);
			Factor pred_factor = factor_graph.addFactor(pred_table, model_params.getTopicTypeVariable(col),
					model_params.getPredicateVariable(col));
			pred_factor.setDirectedTo(model_params.getPredicateVariable(col));
			pred_factor.setLabel("RootPredFactor" + col);
			current_predicate_factors[col] = pred_factor;
			
//			factor_graph.join(current_topic_factor, current_predicate_factors[col]);
			
			for (int row = 0; row < table_info.numRows(); row++) {
				prev_factor = current_entity_factors[row][col];
				if (prev_factor != null) {
					factor_graph.remove(prev_factor);
				}
				IFactorTable entity_factor_table = model_params.getEntityFactor(max_num_cols, col, row);
				Factor entity_factor = factor_graph.addFactor(entity_factor_table, model_params.getTopicVariable(col),
						model_params.getPredicateVariable(col), model_params.getEntityVariable(row, col));
//				if (row == 0)
//					PrettyPrinter.prettyPrintEntityFactor(entity_factor_table, model_params.getTopicVariable(col),
//						model_params.getPredicateVariable(col), model_params.getEntityVariable(row, col), table_info);
				entity_factor.setDirectedTo(model_params.getEntityVariable(row, col));
				entity_factor.setLabel("RootEntFactor" + "_" + col + "_" + row);
				current_entity_factors[row][col] = entity_factor;
			}
		
		}
		
		for (int col = 0; col < backbone.length; col++) {
			if (backbone[col] == -2) {
				break;
			}
			int subj_col = backbone[col];
			if (prev_backbone[col] == subj_col) {
				continue;
			}
			if (subj_col == -1) {
				root_col = col;
			} else {
				for (int row = 0; row < table_info.numRows(); row++) {
					Factor prev_factor = current_entity_factors[row][col];
					if (prev_factor != null) {
						factor_graph.remove(prev_factor);
					}
					IFactorTable entity_factor_table = model_params.getEntityFactor(subj_col, col, row);
					Factor entity_factor = factor_graph.addFactor(entity_factor_table, model_params.getEntityVariable(row, subj_col),
							model_params.getPredicateVariable(col), model_params.getEntityVariable(row, col));
//					PrettyPrinter.prettyPrintEntityFactor(entity_factor_table, model_params.getEntityVariable(row, subj_col),
//							model_params.getPredicateVariable(col), model_params.getEntityVariable(row, col), table_info);
					entity_factor.setDirectedTo(model_params.getEntityVariable(row, col));
					entity_factor.setLabel("EntFactor" + subj_col + "_" + col + "_" + row);
					current_entity_factors[row][col] = entity_factor;
				}
				Factor prev_factor = current_predicate_factors[col];
				if (prev_factor != null) {
					factor_graph.remove(prev_factor);
				}
//				System.out.println(subj_col + " " + col);
				IFactorTable pred_table = model_params.getPredicateFactor(subj_col, col);
				Factor pred_factor = factor_graph.addFactor(pred_table, model_params.getColumnTypeVariable(subj_col),
						model_params.getPredicateVariable(col));
//				PrettyPrinter.prettyPrintPredicateFactor(pred_table, model_params.getColumnTypeVariable(subj_col), 
//						model_params.getPredicateVariable(col), table_info);
				pred_factor.setDirectedTo(model_params.getPredicateVariable(col));
				pred_factor.setLabel("PredFactor" + col);
				current_predicate_factors[col] = pred_factor;
			}
			prev_backbone[col] = subj_col;
		}
//		System.out.println(root_col);

	
		ISchedule schedule = factor_graph.getSchedule();
		Iterator<IScheduleEntry> it = schedule.iterator();
		while (it.hasNext()) {
			IScheduleEntry entry = it.next();
			if (entry instanceof com.analog.lyric.dimple.schedulers.scheduleEntry.NodeScheduleEntry) {
				NodeScheduleEntry node_entry = (NodeScheduleEntry)entry;
//				System.out.println(node_entry.getNode().getLabel());
			}
//			System.out.println(entry.getClass());
		}
//		factor_graph.setSolverFactory(new com.analog.lyric.dimple.solvers.sumproduct.Solver());
//		STableFactor sfactor = (STableFactor)(factor_graph.getSolver());
//		sfactor.setK(5);
		factor_graph.setNumSteps(100);
		factor_graph.solve();
		System.out.println(factor_graph.getBetheFreeEnergy());
		return factor_graph.getBetheFreeEnergy();
	}
	
	private Factor setUpColumnTypeFactor(int col) {
		Pair<int[][], double[]> sparse_info = ColumnTypeFactor.createFactor(model_params.getPredicateVariable(col),
				model_params.getColumnTypeVariable(col), model_params);
		IFactorTable column_type_table = FactorTable.create(sparse_info.getFirst(), sparse_info.getSecond(),
				model_params.getPredicateVariable(col).getDomain(),
				model_params.getColumnTypeVariable(col).getDomain());
		Factor column_type_factor = factor_graph.addFactor(column_type_table,
				model_params.getPredicateVariable(col),
				model_params.getColumnTypeVariable(col));
		column_type_factor.setDirectedTo(model_params.getColumnTypeVariable(col));
		column_type_factor.setLabel("ColType" + (col + 1));
		return column_type_factor;
	}
	
	public void run() {
		double max_score = Double.MAX_VALUE;
		TableInformation table_info = model_params.getTableInformation();
		int num_cols = table_info.numVisibleCols();
		int max_num_cols = model_params.getMaxColumns();
		int[] column_is_leaf = new int[max_num_cols];
		prev_backbone = new int[max_num_cols];
		backbones_checked = new ArrayList<Pair<int[], Double> >();
		for (int c = 0; c < max_num_cols; c++) {
			prev_backbone[c] = -2;
			column_is_leaf[c] = table_info.isLeafColumn(c) ? 1 : 0;
		}
		
		int hstart = 0;
		if (set_backbone != null) {
			hstart = set_backbone.length - table_info.numVisibleCols();
		}
		int num_backbones_checked = 0;
		for (int h = hstart; h < model_params.getMaxNumHiddenColumns() + 1; h++) {
			if (set_backbone != null && set_backbone.length != (num_cols + h)) {
				continue;
			}
			if (h > 0) {
//				setUpColumnTypeFactor(num_cols + h - 1);
			}
			
			ArrayList<int[]> backbones;
			if (set_backbone != null) {
				backbones = new ArrayList<int[]>();
				backbones.add(set_backbone);
			} else if (use_best_backbone) {
				backbones = new ArrayList<int[]>();
				ArrayList<int[]> backbone = BackboneCreator.findGreedyBackbonesOfSizeN(num_cols + h, model_params);
				if (backbone == null || backbone.size() == 0) {  // default to finding all backbones.
					System.out.println("Best greedy backbone could not be found. Defaulting to all backbones.");
					backbones = BackboneCreator.findAllBackbonesOfSizeN(num_cols + h, column_is_leaf);
				} else {
					backbones.addAll(backbone);
				}
			} else {
				backbones = BackboneCreator.findAllBackbonesOfSizeN(num_cols + h, column_is_leaf);
			}
			int num_backbones = backbones.size();
			for (int b = 0; b < num_backbones; b++) {
				if (b % 100 == 0)
					System.out.println(b);
				int[] backbone = backbones.get(b);
				for (int i = 0; i < backbone.length; i++) {
					System.out.print(backbone[i] + " ");
				}
				System.out.println();
			}
			System.out.println("num backbones: " + num_backbones);
			
			for (int b = 0; b < num_backbones; b++) {
				if (b % 100 == 0)
					System.out.println(b);
				int[] backbone = backbones.get(b);
				for (int i = 0; i < backbone.length; i++) {
					System.out.print(backbone[i] + " ");
				}
				System.out.println();
				if (!ValidBackbone(backbone)) {
					continue;
				}
				double score = runInferenceForBackbone(backbone);
				backbones_checked.add(new Pair<int[], Double>(backbone, score));
				System.out.println(score);
				if (score < max_score) {
					max_score = score;
					storeBestBeliefs(backbone);
				}
				num_backbones_checked++;
			}
		}
		System.out.println("num backbones checked: " + num_backbones_checked);
		for (int i = 0; i < best_backbone.length; i++) {
			System.out.print(best_backbone[i] + " ");
		}
		System.out.println();
		best_bethe_free_energy = max_score;
//		for (int r = 0; r < this.best_entity_vars_beliefs.length; r++) {
//			for (int c = 0; c < this.best_entity_vars_beliefs[r].length; c++) {
//				System.out.print(this.best_entity_vars_beliefs[r][c]);
//			}
//			System.out.println();
//		}
//		for (int col = 0; col < best_backbone.length; col++) {
//			Discrete type_var = model_params.getColumnTypeVariable(col);
//			Object[] types = type_var.getDomain().getElements();
//			double[] input = type_var.getInput();
//			System.out.println(col + ":");
//			for (int t = 0; t < types.length; t++) {
//				System.out.println("  " + table_info.getStringForId((Integer)(types[t])) + " " + best_col_type_vars_beliefs[col][t] + " " + input[t]);
//			}
//		}
		PrettyPrinter.prettyPrintTableInterpretation(this, this.best_backbone, this.best_entity_vars_beliefs,
				this.best_predicates_vars_beliefs, model_params.getTopicVariable(this.best_topic_var_idx), this.best_topic_var_beliefs, 
				model_params.getTopicTypeVariable(this.best_topic_var_idx), this.best_topic_type_var_beliefs, 5);
	}
	
	private boolean ValidBackbone(int[] backbone) {
		int num_cols = model_params.getTableInformation().numVisibleCols();
		int last_par = -2;
		for (int i = num_cols; i < backbone.length; i++) {
			if (backbone[i] >= num_cols) {
				if (VERBOSE) {
					System.out.println("CVT column subject of another CVT column.");
				}
				return false;
			}
			if (backbone[i] < last_par) {
				if (VERBOSE) {
					System.out.println("Symmetric to another backbone.");
				}
				return false;
			}
			last_par = backbone[i];
			
		}
		for (int i = 0; i < backbone.length; i++) {
			if (backbone[i] == -1) {
				continue;
			}
			if (backbone[i] == 0) {
				return true;
			}
			int subj_col = backbone[i];
			if (subj_col > num_cols) {
				subj_col = num_cols;
			}
//			if (!model_params.isValidBackboneConnection(backbone[i], i)) {
//				System.out.println("!");
//				return false;
//			}
		}
		return true;
	}

	private void storeBestBeliefs(int[] backbone) {
		best_backbone = backbone.clone();
		int backbone_size = ArrayUtils.indexOf(best_backbone, -2);
		if (backbone_size == -1) {
			backbone_size = best_backbone.length;
		}
		for (int col = 0; col < backbone_size; col++) {
			best_predicates_vars_beliefs[col] = model_params.getPredicateVariable(col).getBelief().clone();
			best_col_type_vars_beliefs[col] = model_params.getColumnTypeVariable(col).getBelief().clone();
			for (int row = 0; row < model_params.getTableInformation().numRows(); row++) {
				best_entity_vars_beliefs[row][col] = model_params.getEntityVariable(row, col).getBelief().clone();
			}
		}
		int root_col = ArrayUtils.indexOf(backbone, -1);
		best_topic_var_beliefs = model_params.getTopicVariable(root_col).getBelief().clone();
		best_topic_type_var_beliefs = model_params.getTopicTypeVariable(root_col).getBelief().clone();
		best_topic_var_idx = root_col;
		best_backbone_length = backbone_size;
	}

	public static void main(String[] args) {
		try {
			long start_setup = System.nanoTime();
			TableInferer table_inference = new TableInferer();
			if (args.length == 6)
				table_inference.setup(args[0], args[1], args[2], args[3], args[4], args[5]);
			else
				table_inference.setup(args[0], args[1], args[2], args[3], args[4], "");
	//		if (true)
	//			return;
			long end_setup = System.nanoTime();
			System.out.println((end_setup - start_setup)/ ((double)1000000000));
			long start_inf = System.nanoTime();
			table_inference.run();
			long end_inf = System.nanoTime();
			System.out.println((end_inf - start_inf)/ ((double)1000000000));	
		} catch (IOException ex) {
			ex.printStackTrace();
		}
	}
}
